【精选】Pytorch网络模型权重初始化、保存与加载模型、加载预训练模型、按需设置学习率 |
您所在的位置:网站首页 › pytorch 加载部分权重 › 【精选】Pytorch网络模型权重初始化、保存与加载模型、加载预训练模型、按需设置学习率 |
前言
在我们对神经网络模型进行训练时,往往需要对模型进行初始化或者加载预训练模型。本文将对模型的权重初始化与加载预训练模型做一个学习记录,以便后续查询使用。 权重初始化 常见的初始化方法PyTorch 在 torch.nn.init 中提供了常用的初始化方法函数,这里主要简要介绍Xavier初始化与kaiming初始化。 Xavier初始化Xavier 初始化方法,方法来源于2010年的一篇论文《Understanding the difficulty of training deep feedforward neural networks》 公式推导是从“方差一致性”出发,初始化的分布有均匀分布和正态分布两种。 Xavier 均匀分布torch.nn.init.xavier_uniform_(tensor, gain=1.0) 该初始化方法服从均匀分布 U ∼ ( − a , a ) U\sim(-a,a) U∼(−a,a),其中a为: a = gain × 6 fan_in + fan_out a=\operatorname{gain} \times \sqrt{\frac{6}{\operatorname{fan\_in} +\operatorname{fan\_out}}} a=gain×fan_in+fan_out6 该初始化方法中有一个参数 gain,增益的大小是依据激活函数类型来设定 eg: nn.init.xavier_uniform_(w, gain=nn.init.calculate_gain(‘relu’)) PS:上述初始化方法,也称为 Glorot initialization 使用方法示例: for m in model.modules(): if isinstance(m, (nn.Conv2d, nn.Linear)): nn.init.xavier_uniform_(m.weight) Xavier正态分布torch.nn.init.xavier_normal_(tensor, gain=1.0) 该初始化方法服从正态分布 N ( 0 , s t d 2 ) \mathcal{N}\left(0, \mathrm{std}^{2}\right) N(0,std2): std = gain × 2 fan_in + fan_out \operatorname{std}=\operatorname{gain} \times \sqrt{\frac{2}{\operatorname{fan\_in} +\operatorname{fan\_out}}} std=gain×fan_in+fan_out2 使用方法示例: for m in model.modules(): if isinstance(m, (nn.Conv2d, nn.Linear)): nn.init.xavier_normal_(m.weight) kaiming初始化kaiming初始化,方法来源于2015年的一篇论文《 Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification》 公式推导同样从“方差一致性”出法,kaiming是针对xavier初始化方法在relu这一类激活函数表现不佳而提出的改进,详细可以参看论文。 kaiming均匀分布torch.nn.init.kaiming_uniform_(tensor, a=0, mode=‘fan_in’, nonlinearity=‘leaky_relu’) 该初始化方法服从均匀分布 U ( \mathcal{U}( U( -bound, bound ) ) ): bound = gain × 3 fan_mode \text { bound }=\operatorname{gain} \times \sqrt{\frac{3}{\text { fan\_mode }}} bound =gain× fan_mode 3 其中,a为激活函数的负半轴的斜率,mode可选为fan_in 或fan_out, fan_in使正向传播时,方差一致; fan_out使反向传播时,方差一致。 nonlinearity 建议选择 relu 和 leaky_relu ,默认值为 leaky_relu kaiming正态分布torch.nn.init.kaiming_normal_(tensor, a=0, mode=‘fan_in’, nonlinearity=‘leaky_relu’) 该初始化方法服从正态分布 N ( 0 , s t d 2 ) \mathcal{N}\left(0, \mathrm{std}^{2}\right) N(0,std2): s t d = gain fan_mode \mathrm{std}=\frac{\text { gain }}{\sqrt{\text { fan\_mode }}} std= fan_mode gain 其中,a为激活函数的负半轴的斜率,mode可选为fan_in 或fan_out, fan_in使正向传播时,方差一致; fan_out使反向传播时,方差一致。 nonlinearity 建议选择 relu 和 leaky_relu ,默认值为 leaky_relu 模型权重初始化 # 定义权值初始化 def initialize_weights(self): for m in self.modules(): if isinstance(m, nn.Conv2d): torch.nn.init.xavier_normal_(m.weight.data) if m.bias is not None: m.bias.data.zero_() elif isinstance(m, nn.BatchNorm2d): m.weight.data.fill_(1) m.bias.data.zero_() elif isinstance(m, nn.Linear): torch.nn.init.normal_(m.weight.data, 0, 0.01) m.bias.data.zero_() 保存与加载模型pytorch在保存模型时,可以保存整个神经网络的的结构信息和模型参数信息,save的对象是网络net;也可以只保存神经网络的训练模型参数,save的对象是net.state_dict()。 # 保存和加载整个模型 torch.save(model_object, 'model.pth') model = torch.load('model.pth') # 仅保存和加载模型参数 torch.save(model_object.state_dict(), 'params.pth') model_object.load_state_dict(torch.load('params.pth')) 加载预训练模型 # load params, 这里加载的是模型的参数,不是整个模型 pretrained_dict = torch.load('net_params.pkl') # 仅保存了整个模型, 需要使用以下语句 # pretrained_dict = torch.load('net_params.pkl').state_dict() # 获取当前网络的dict net_state_dict = net.state_dict() # 剔除不匹配的权值参数 pretrained_dict_1 = {k: v for k, v in pretrained_dict.items() if k in net_state_dict} # 更新新模型参数字典 net_state_dict.update(pretrained_dict_1) # 将包含预训练模型参数的字典"放"到新模型中 net.load_state_dict(net_state_dict) pytorch预训练模型的简单修改与使用以resnet预训练模型举例,resnet源代码的pytorch官方实现。 resnet网络最后一层分类层fc是对1000种类型进行划分,如果自己的数据集只有6类,可以只对fc层进行修改: #调用模型 model = torchvision.models.resnet50(pretrained=True) #提取fc层中固定的参数 fc_features = model.fc.in_features #修改类别 model.fc = nn.Linear(fc_features, 6) 按需设置学习率 # ================================= # # 按需设置学习率 # ================================= # # 将fc3层的参数从原始网络参数中剔除 ignored_params = list(map(id, net.fc3.parameters())) base_params = filter(lambda p: id(p) not in ignored_params, net.parameters()) # 为fc3层设置需要的学习率 optimizer = optim.SGD([ {'params': base_params}, {'params': net.fc3.parameters(), 'lr': lr_init*10}], lr_init, momentum=0.9, weight_decay=1e-4) criterion = nn.CrossEntropyLoss() # 选择损失函数 scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.1) # 设置学习率下降策略 参考资料 《Pytorch模型训练实用教程》《Understanding the difficulty of training deep feedforward neural networks》《 Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification》PyTorch Documentation |
今日新闻 |
推荐新闻 |
CopyRight 2018-2019 办公设备维修网 版权所有 豫ICP备15022753号-3 |